package com.atlin.project.config;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;

public class CountingHttpServletResponseWrapper extends HttpServletResponseWrapper {
    private CountingServletOutputStream countingStream;
    private PrintWriter writer;

    public CountingHttpServletResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if (countingStream == null) {
            countingStream = new CountingServletOutputStream(super.getOutputStream());
        }
        return countingStream;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (writer == null) {
            writer = new PrintWriter(new OutputStreamWriter(getOutputStream(), getCharacterEncoding()));
        }
        return writer;
    }

    public long getByteCount() {
        return countingStream == null ? 0 : countingStream.getByteCount();
    }
}

class CountingServletOutputStream extends ServletOutputStream {
    private final ServletOutputStream outputStream;
    private long byteCount = 0;

    public CountingServletOutputStream(ServletOutputStream outputStream) {
        this.outputStream = outputStream;
    }

    @Override
    public void write(int b) throws IOException {
        outputStream.write(b);
        byteCount++;
    }

    @Override
    public void write(byte[] b, int off, int len) throws IOException {
        outputStream.write(b, off, len);
        byteCount += len;
    }

    public long getByteCount() {
        return byteCount;
    }

    @Override
    public boolean isReady() {
        return outputStream.isReady();
    }

    @Override
    public void setWriteListener(WriteListener writeListener) {
        outputStream.setWriteListener(writeListener);
    }
}
